"""
Prompt Manager for ComfyUI-xhh - MySQL版本
集成promptManager-server功能到ComfyUI中
"""

import os
import json
import pymysql
import base64
import urllib.parse
from datetime import datetime
from typing import List, Dict, Optional
import threading
import time

class PromptManager:
    """提示词管理器 - MySQL版本"""
    
    def __init__(self, mysql_config: dict = None):
        if mysql_config is None:
            self.mysql_config = {
                'host': 'localhost',
                'port': 3306,
                'user': 'root',
                'password': 'root',
                'database': 'comfyui-xhh',
                'charset': 'utf8mb4'
            }
        else:
            self.mysql_config = mysql_config
        
        self.init_database()
    
    def get_connection(self):
        """获取MySQL数据库连接"""
        try:
            return pymysql.connect(**self.mysql_config)
        except Exception as e:
            print(f"[PromptManager] MySQL连接失败: {e}")
            return None
    
    def init_database(self):
        """初始化数据库"""
        try:
            conn = self.get_connection()
            if not conn:
                print("[PromptManager] 无法连接到MySQL数据库")
                return
            
            with conn.cursor() as cursor:
                # 创建prompts表（去掉图片相关字段）
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS prompts (
                        id INT AUTO_INCREMENT PRIMARY KEY,
                        chrome_id VARCHAR(255) UNIQUE,
                        text TEXT NOT NULL,
                        source VARCHAR(50) DEFAULT 'workflow',
                        is_favorite BOOLEAN DEFAULT FALSE,
                        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                        last_synced DATETIME DEFAULT CURRENT_TIMESTAMP
                    )
                ''')
                
                # 创建images表（增加缩略图路径字段）
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS images (
                        id INT AUTO_INCREMENT PRIMARY KEY,
                        filename VARCHAR(255) NOT NULL,
                        file_path TEXT NOT NULL,
                        thumbnail_path TEXT,
                        file_size BIGINT,
                        width INT,
                        height INT,
                        mime_type VARCHAR(100),
                        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
                    )
                ''')
                
                # 创建prompt_images关系表
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS prompt_images (
                        id INT AUTO_INCREMENT PRIMARY KEY,
                        prompt_id INT NOT NULL,
                        image_id INT NOT NULL,
                        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE,
                        FOREIGN KEY (image_id) REFERENCES images(id) ON DELETE CASCADE,
                        UNIQUE KEY unique_prompt_image (prompt_id, image_id)
                    )
                ''')
                
                # 创建loras表
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS loras (
                        id VARCHAR(255) PRIMARY KEY,
                        name VARCHAR(255) NOT NULL,
                        path TEXT NOT NULL,
                        size BIGINT DEFAULT 0,
                        tags TEXT,
                        description TEXT,
                        `trigger` TEXT,
                        md5 VARCHAR(32),
                        is_favorite BOOLEAN DEFAULT FALSE,
                        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
                    )
                ''')
                
                # 创建lora_images关系表
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS lora_images (
                        id INT AUTO_INCREMENT PRIMARY KEY,
                        lora_id VARCHAR(255) NOT NULL,
                        image_id INT NOT NULL,
                        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        FOREIGN KEY (lora_id) REFERENCES loras(id) ON DELETE CASCADE,
                        FOREIGN KEY (image_id) REFERENCES images(id) ON DELETE CASCADE,
                        UNIQUE KEY unique_lora_image (lora_id, image_id)
                    )
                ''')
                
                # 创建categories表
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS categories (
                        id INT AUTO_INCREMENT PRIMARY KEY,
                        category_key VARCHAR(255) NOT NULL,
                        category_name VARCHAR(255) NOT NULL,
                        description TEXT,
                        parent_id INT,
                        level INT DEFAULT 0,
                        data TEXT,
                        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                        UNIQUE KEY unique_category_key (category_key)
                    )
                ''')
                
                conn.commit()
                
                # 执行数据库迁移
                self._migrate_loras_table(cursor)
                
                print("[PromptManager] 数据库初始化完成")
            
        except Exception as e:
            print(f"[PromptManager] 数据库初始化失败: {e}")
        finally:
            if conn:
                conn.close()
    
    def add_prompt(self, text: str, source: str = 'workflow', chrome_id: str = None, is_favorite: bool = False) -> int:
        """添加提示词"""
        try:
            conn = self.get_connection()
            if not conn:
                return None
            
            with conn.cursor() as cursor:
                if chrome_id is None:
                    chrome_id = str(int(time.time() * 1000))
                
                cursor.execute('''
                    INSERT INTO prompts (chrome_id, text, source, is_favorite)
                    VALUES (%s, %s, %s, %s)
                ''', (chrome_id, text, source, is_favorite))
                
                prompt_id = cursor.lastrowid
                conn.commit()
                
                print(f"[PromptManager] 添加提示词成功: {prompt_id}")
                return prompt_id
                
        except Exception as e:
            print(f"[PromptManager] 添加提示词失败: {e}")
            return None
        finally:
            if conn:
                conn.close()
    
    def get_prompts(self, limit: int = 100, offset: int = 0, 
                   search: str = None, favorite_only: bool = False) -> List[Dict]:
        """获取提示词列表"""
        try:
            conn = self.get_connection()
            if not conn:
                return []
            
            with conn.cursor() as cursor:
                query = "SELECT * FROM prompts WHERE 1=1"
                params = []
                
                if search:
                    query += " AND text LIKE %s"
                    params.append(f'%{search}%')
                
                if favorite_only:
                    query += " AND is_favorite = TRUE"
                
                query += " ORDER BY created_at DESC LIMIT %s OFFSET %s"
                params.extend([limit, offset])
                
                cursor.execute(query, params)
                columns = [desc[0] for desc in cursor.description]
                results = []
                
                for row in cursor.fetchall():
                    prompt = dict(zip(columns, row))
                    # 转换datetime对象为字符串
                    for key, value in prompt.items():
                        if isinstance(value, datetime):
                            prompt[key] = value.isoformat()
                    results.append(prompt)
                
                return results
                
        except Exception as e:
            print(f"[PromptManager] 获取提示词列表失败: {e}")
            return []
        finally:
            if conn:
                conn.close()
    
    def get_prompts_count(self, search: str = None, favorite_only: bool = False) -> int:
        """获取提示词总数"""
        try:
            conn = self.get_connection()
            if not conn:
                return 0
            
            with conn.cursor() as cursor:
                query = "SELECT COUNT(*) FROM prompts WHERE 1=1"
                params = []
                
                if search:
                    query += " AND text LIKE %s"
                    params.append(f'%{search}%')
                
                if favorite_only:
                    query += " AND is_favorite = TRUE"
                
                cursor.execute(query, params)
                result = cursor.fetchone()
                return result[0] if result else 0
                
        except Exception as e:
            print(f"[PromptManager] 获取提示词总数失败: {e}")
            return 0
        finally:
            if conn:
                conn.close()
    
    def get_prompt_by_id(self, prompt_id: int) -> Optional[Dict]:
        """根据ID获取提示词"""
        try:
            conn = self.get_connection()
            if not conn:
                return None
            
            with conn.cursor() as cursor:
                cursor.execute("SELECT * FROM prompts WHERE id = %s", (prompt_id,))
                result = cursor.fetchone()
                
                if result:
                    columns = [desc[0] for desc in cursor.description]
                    prompt = dict(zip(columns, result))
                    # 转换datetime对象为字符串
                    for key, value in prompt.items():
                        if isinstance(value, datetime):
                            prompt[key] = value.isoformat()
                    return prompt
                
                return None
                
        except Exception as e:
            print(f"[PromptManager] 获取提示词失败: {e}")
            return None
        finally:
            if conn:
                conn.close()
    
    def update_prompt(self, prompt_id: int, **kwargs) -> bool:
        """更新提示词"""
        try:
            conn = self.get_connection()
            if not conn:
                return False
            
            with conn.cursor() as cursor:
                # 构建更新语句
                set_clauses = []
                params = []
                
                for key, value in kwargs.items():
                    if key in ['text', 'image_data', 'original_image', 'source', 'is_favorite']:
                        set_clauses.append(f"{key} = %s")
                        params.append(value)
                
                if not set_clauses:
                    return False
                
                params.append(prompt_id)
                query = f"UPDATE prompts SET {', '.join(set_clauses)} WHERE id = %s"
                
                cursor.execute(query, params)
                conn.commit()
                
                return cursor.rowcount > 0
                
        except Exception as e:
            print(f"[PromptManager] 更新提示词失败: {e}")
            return False
        finally:
            if conn:
                conn.close()
    
    def delete_prompt(self, prompt_id: int) -> bool:
        """删除提示词及其关联的图片关系"""
        try:
            conn = self.get_connection()
            if not conn:
                return False
            
            with conn.cursor() as cursor:
                # 先删除关联的图片关系
                cursor.execute("DELETE FROM prompt_images WHERE prompt_id = %s", (prompt_id,))
                print(f"[PromptManager] 删除了 {cursor.rowcount} 条图片关联记录")
                
                # 再删除提示词
                cursor.execute("DELETE FROM prompts WHERE id = %s", (prompt_id,))
                deleted_count = cursor.rowcount
                conn.commit()
                
                print(f"[PromptManager] 删除了 {deleted_count} 条提示词记录")
                return deleted_count > 0
                
        except Exception as e:
            print(f"[PromptManager] 删除提示词失败: {e}")
            return False
        finally:
            if conn:
                conn.close()
    
    def delete_image_record(self, image_id: int) -> bool:
        """删除图片记录及其关联关系"""
        try:
            conn = self.get_connection()
            if not conn:
                return False
            
            with conn.cursor() as cursor:
                # 先删除关联关系
                cursor.execute("DELETE FROM prompt_images WHERE image_id = %s", (image_id,))
                print(f"[PromptManager] 删除了 {cursor.rowcount} 条图片关联记录")
                
                # 再删除图片记录
                cursor.execute("DELETE FROM images WHERE id = %s", (image_id,))
                deleted_count = cursor.rowcount
                conn.commit()
                
                print(f"[PromptManager] 删除了 {deleted_count} 条图片记录")
                return deleted_count > 0
                
        except Exception as e:
            print(f"[PromptManager] 删除图片记录失败: {e}")
            return False
        finally:
            if conn:
                conn.close()
    
    def toggle_favorite(self, prompt_id: int) -> bool:
        """切换收藏状态"""
        try:
            conn = self.get_connection()
            if not conn:
                return False
            
            with conn.cursor() as cursor:
                # 获取当前收藏状态
                cursor.execute("SELECT is_favorite FROM prompts WHERE id = %s", (prompt_id,))
                result = cursor.fetchone()
                
                if result:
                    current_favorite = result[0]
                    new_favorite = not current_favorite
                    
                    cursor.execute("UPDATE prompts SET is_favorite = %s WHERE id = %s", 
                                 (new_favorite, prompt_id))
                    conn.commit()
                    
                    return True
                
                return False
                
        except Exception as e:
            print(f"[PromptManager] 切换收藏状态失败: {e}")
            return False
        finally:
            if conn:
                conn.close()
    
    def get_prompt_count(self, search: str = None, favorite_only: bool = False) -> int:
        """获取提示词总数"""
        try:
            conn = self.get_connection()
            if not conn:
                return 0
            
            with conn.cursor() as cursor:
                query = "SELECT COUNT(*) FROM prompts WHERE 1=1"
                params = []
                
                if search:
                    query += " AND text LIKE %s"
                    params.append(f'%{search}%')
                
                if favorite_only:
                    query += " AND is_favorite = TRUE"
                
                cursor.execute(query, params)
                result = cursor.fetchone()
                
                return result[0] if result else 0
                
        except Exception as e:
            print(f"[PromptManager] 获取提示词总数失败: {e}")
            return 0
        finally:
            if conn:
                conn.close()
    
    def get_prompt_images(self, prompt_id: int) -> List[Dict]:
        """获取提示词关联的图片"""
        try:
            conn = self.get_connection()
            if not conn:
                return []
            
            with conn.cursor() as cursor:
                cursor.execute('''
                    SELECT i.id, i.filename, i.file_path, i.file_size, i.width, i.height, i.mime_type
                    FROM images i
                    JOIN prompt_images pi ON i.id = pi.image_id
                    WHERE pi.prompt_id = %s
                    ORDER BY pi.created_at DESC
                ''', (prompt_id,))
                
                columns = [desc[0] for desc in cursor.description]
                results = []
                
                for row in cursor.fetchall():
                    image = dict(zip(columns, row))
                    results.append(image)
                
                return results
                
        except Exception as e:
            print(f"[PromptManager] 获取提示词图片失败: {e}")
            return []
        finally:
            if conn:
                conn.close()
    
    def add_image(self, filename: str, file_path: str, file_size: int = None, 
                  width: int = None, height: int = None, mime_type: str = None) -> int:
        """添加图片记录"""
        try:
            conn = self.get_connection()
            if not conn:
                return None
            
            with conn.cursor() as cursor:
                cursor.execute('''
                    INSERT INTO images (filename, file_path, file_size, width, height, mime_type)
                    VALUES (%s, %s, %s, %s, %s, %s)
                ''', (filename, file_path, file_size, width, height, mime_type))
                
                image_id = cursor.lastrowid
                conn.commit()
                
                print(f"[PromptManager] 添加图片记录成功: {image_id}")
                return image_id
                
        except Exception as e:
            print(f"[PromptManager] 添加图片记录失败: {e}")
            return None
        finally:
            if conn:
                conn.close()
    
    def add_prompt_image_relation(self, prompt_id: int, image_id: int) -> bool:
        """添加提示词图片关联"""
        try:
            conn = self.get_connection()
            if not conn:
                return False
            
            with conn.cursor() as cursor:
                # 检查关联是否已存在
                cursor.execute('''
                    SELECT id FROM prompt_images 
                    WHERE prompt_id = %s AND image_id = %s
                ''', (prompt_id, image_id))
                
                if cursor.fetchone():
                    return True  # 关联已存在
                
                # 创建新关联
                cursor.execute('''
                    INSERT INTO prompt_images (prompt_id, image_id)
                    VALUES (%s, %s)
                ''', (prompt_id, image_id))
                
                conn.commit()
                return True
                
        except Exception as e:
            print(f"[PromptManager] 添加提示词图片关联失败: {e}")
            return False
        finally:
            if conn:
                conn.close()
    
    def _migrate_loras_table(self, cursor):
        """迁移loras表结构"""
        try:
            print("[PromptManager] 开始迁移loras表...")
            
            # 检查是否需要添加md5字段
            cursor.execute("SHOW COLUMNS FROM loras LIKE 'md5'")
            if not cursor.fetchone():
                print("[PromptManager] 添加md5字段...")
                cursor.execute("ALTER TABLE loras ADD COLUMN md5 VARCHAR(32) AFTER `trigger`")
            
            # 检查是否需要删除prompt字段
            cursor.execute("SHOW COLUMNS FROM loras LIKE 'prompt'")
            if cursor.fetchone():
                print("[PromptManager] 删除prompt字段...")
                cursor.execute("ALTER TABLE loras DROP COLUMN prompt")
            
            print("[PromptManager] loras表迁移完成")
            
        except Exception as e:
            print(f"[PromptManager] loras表迁移失败: {e}")
            import traceback
            traceback.print_exc()
    
    def get_prompt_images(self, prompt_id):
        """获取提示词关联的图片列表"""
        try:
            conn = pymysql.connect(**self.mysql_config)
            with conn.cursor() as cursor:
                cursor.execute("""
                    SELECT i.id, i.filename, i.file_path, i.thumbnail_path, i.file_size, i.width, i.height, i.mime_type
                    FROM images i
                    JOIN prompt_images pi ON i.id = pi.image_id
                    WHERE pi.prompt_id = %s
                    ORDER BY pi.created_at DESC
                """, (prompt_id,))
                
                images = []
                for row in cursor.fetchall():
                    image = {
                        "id": row[0],
                        "filename": row[1],
                        "file_path": row[2],
                        "thumbnail_path": row[3],
                        "file_size": row[4],
                        "width": row[5],
                        "height": row[6],
                        "mime_type": row[7]
                    }
                    images.append(image)
                
                return images
                
        except Exception as e:
            print(f"[PromptManager] 获取提示词图片失败: {e}")
            return []
        finally:
            if conn:
                conn.close()
